//
// Created by vivi on 09/02/2018.
//

#include "core/collider/sdfcollider.h"
#include "core/util/stringutil.h"
#include "core/util/fileutil.h"
#include "core/util/mathutil.h"
#include "core/util/eigenutil.h"
#include <fstream>
#include <iostream>
#include <algorithm>
#include <memory>

using namespace std;
using namespace Eigen;

namespace HairEngine {
    SDFColliderCell::SDFColliderCell(const SDFColliderCell &other) = default;

	SDFCollider::SDFCollider(const std::string &filePath) {
		if (!StringUtility::endswith(filePath, ".sdf") && !StringUtility::endswith(filePath, ".sdf2"))
			Error("The correct file format should be endsed with .sdf");

		fstream sdfFile(filePath, ios::in | ios::binary);
		if (!sdfFile.is_open())
			Error("Unexpected result while opening file \"%s\"", filePath.c_str());

		Info("Begin to read sdf file of \"%s\"\n", filePath.c_str());

		nx = (size_t)FileUtility::binaryReadInt32(sdfFile);
		Assert(sdfFile);
		ny = (size_t)FileUtility::binaryReadInt32(sdfFile);
		Assert(sdfFile);
		nz = (size_t)FileUtility::binaryReadInt32(sdfFile);
		Assert(sdfFile);
		nyz = ny * nz;
		vertexCount = nx * nyz;

		minP = FileUtility::binaryReadVector3f(sdfFile);
		Assert(sdfFile);
		if (StringUtility::endswith(filePath, ".sdf")) {
			//.sdf
			//The length of voxel in .sdf are equal in each dimensions
			const float tmp = FileUtility::binaryReadReal32(sdfFile);
			cellSize << tmp, tmp, tmp;
		}
		else {
			//.sdf2
			cellSize = FileUtility::binaryReadVector3f(sdfFile);
		}
		cellSizeReciprocal << 1.0f / cellSize(0)
				, 1.0f / cellSize(1)
				, 1.0f / cellSize(2);
		Assert(sdfFile);
		maxP << minP(0) + (nx - 1) * cellSize(0)
				, minP(1) + (ny - 1) * cellSize(1)
				, minP(2) + (nz - 1) * cellSize(2);

		bounding = Eigen::AlignedBox3f(minP, maxP);

		vertexes = new SDFColliderVertex[vertexCount];
		for (size_t i = 0; i < vertexCount; ++i)
			vertexes[i].distance = FileUtility::binaryReadReal32(sdfFile);

		Info("Done reading sdf file of \"%s\"\n", filePath.c_str());
		Info("Information about sdf file \"%s\"\n", filePath.c_str());
		Info("Number of vertexes: nx = %d, ny = %d, nz = %d, nyz = %d, cellSize = %s\n", (int)nx, (int)ny, (int)nz, (int)nyz, EigenUtility::toString(cellSize).c_str());
		Info("Bounding box: minP = %s, maxP = %s\n", EigenUtility::toString(minP).c_str(), EigenUtility::toString(maxP).c_str());
		Info("The distance of first %d voxel distance are: \n", std::min((int)vertexCount, 1000));
		IF_ON_DEBUG {
			for (size_t i = 0; i < std::min(vertexCount, (size_t)1000); ++i) {
				Info("%f  ", vertexes[i].distance);
				if (i % 1000 == 0)
					Info("\n");
			}
			Info("\n");

			size_t positiveCounter = 0, negativeCounter = 0;
			for (size_t i = 0; i < vertexCount; ++i) {
				if (vertexes[i].distance > 0)
					++positiveCounter;
				else
					++negativeCounter;
			}

			Info("Postive counter = %d, negative count = %d\n", (int)positiveCounter, (int)negativeCounter);
		};
	}

	SDFColliderCell SDFCollider::getCell(const Eigen::Vector3f &pos, Eigen::Vector3f *outT) const {
		const Vector3f indexInFloat = (pos - minP).cwiseProduct(cellSizeReciprocal);

		//Only check the boundary in debug mode
		Assert(MathUtility::between(pos(0), minP(0), maxP(0)));
		Assert(MathUtility::between(pos(1), minP(1), maxP(1)));
		Assert(MathUtility::between(pos(2), minP(2), maxP(2)));

		const auto dx = (size_t)(indexInFloat(0));
		const auto dy = (size_t)(indexInFloat(1));
		const auto dz = (size_t)(indexInFloat(2));

		Cell cell;
		cell.v[0] = vertexes + (dx * nyz + dy * nz + dz);
		cell.v[1] = cell.v[0] + 1;
		cell.v[2] = cell.v[0] + nz;
		cell.v[3] = cell.v[2] + 1;
		cell.v[4] = cell.v[0] + nyz;
		cell.v[5] = cell.v[4] + 1;
		cell.v[6] = cell.v[4] + nz;
		cell.v[7] = cell.v[6] + 1;

		if (outT)
			(*outT) << indexInFloat(0) - dx, indexInFloat(1) - dy, indexInFloat(2) - dz;

		return cell;
	};

	float SDFCollider::distance(const Eigen::Vector3f &pos, Eigen::Vector3f *outGradientPtr) const {

		//If not in the bounding box, we cannot get the distance
		if (!bounding.contains(pos))
			return SDFCollider::SDF_INFINITY;

		//Get the surrounding
		Vector3f t;
		auto cell = getCell(pos, &t);
		const float & tx = t(0);
		const float & ty = t(1);
		const float & tz = t(2);

		//Those are used for prev caculation storage for computation distance and gradient using trilinear interpolation
		float chunkx[2] = {1 - tx, tx};
		float chunky[2] = {1 - ty, ty};
		float chunkz[2] = {1 - tz, tz};
		float chunkxForGradient[2] = { -1.0f / (1 - tx), 1.0f / tx};
		float chunkyForGradient[2] = { -1.0f / (1 - ty), 1.0f / ty};
		float chunkzForGradient[2] = { -1.0f / (1 - tz), 1.0f / tz};

		/*
		 * For distance, the "distance" in the voxel is interpolated with trilinear interpolation,
		 * so if the 8 neighbors are denoted by v000 (the lowest point), v001 (increased z), v010 (increased y),
		 * v011 (increased y and z), v100 (increased x), ... , v111. For the given tx, ty, tz ( 0 < x, y, z < 1), we have
		 *
		 * distanceInXYZ =
		 *         v000 * (1 - tx) * (1 - ty) * (1 - tz) +
		 *         v001 * (1 - tx) * (1 - ty) * tz +
		 *         v010 * (1 - tx) * ty * (1 - tz)
		 *         ... +
		 *         v111 * tx * ty * tz
		 *
		 *  We compute the gradient along with the computation of distance. When summing the corresponding 8 terms
		 *  (described above), if the gradient "z" term is (1 - z), it will donate an v000 * (1 - x) * (1 - y) * -1.0f to the
		 *  gradient "z" term, but the distance donation is v000 * (1 - x) * (1 - y) * (1 - z), we could use the distance donation,
		 *  multiply it by chunkZForGradient[0] and gets the result.
		 */

		//Compute the distance and the gradient
		float distance = 0.0f;
		Vector3f gradient = Vector3f::Zero();

		for (size_t i = 0; i < 8; ++i) {
			size_t boolx = (i >> 2), booly = (i >> 1) & 1, boolz = i & 1;

			float term = cell.v[i]->distance * chunkx[boolx] * chunky[booly] * chunkz[boolz];
			if (outGradientPtr != nullptr) {
				gradient(0) += (boolx ? 1.0f : -1.0f) * cell.v[i]->distance * chunky[booly] * chunkz[boolz];
				gradient(1) += (booly ? 1.0f : -1.0f) * cell.v[i]->distance * chunkx[boolx] * chunkz[boolz];
				gradient(2) += (boolz ? 1.0f : -1.0f) * cell.v[i]->distance * chunkx[boolx] * chunky[booly];
			}
			distance += term;
		}

		if (outGradientPtr != nullptr) {
			/*
			 * Since x = CONSTANT + tx * cellSize.x
			 * d(tx) / dx = 1.0f / cellSize.x
			 * By using the chain rule of derivative, we must multiply the result graident with the reciprocal of cell size
			 * in each dimension.
            */
			*outGradientPtr = gradient.cwiseProduct(cellSizeReciprocal);
		}

		return distance;
	}

	bool SDFCollider::debugWriteZeroIsoContourToFile(const std::string &filePath) {
		fstream contourFile(filePath, ios::out);
		if (!contourFile.is_open()) {
			printf("Error while opening \"%s\"\n", filePath.c_str());
			return false;
		}

		contourFile << "wireframebox"
		            << ' ' << minP(0) << ' ' << minP(1) << ' ' << minP(2)
		            << ' ' << maxP(0) << ' ' << maxP(1) << ' ' << maxP(2) << '\n';
		Assert(contourFile);

		size_t successCounter = 0;
		for (size_t dx = 0; dx < nx - 1; ++dx)
			for (size_t dy = 0; dy < ny - 1; ++dy)
				for (size_t dz = 0; dz < nz - 1; ++dz) {
					Vector3f centerPoint;
					centerPoint << minP(0) + (dx + 0.5f) * cellSize(0)
					            , minP(1) + (dy + 0.5f) * cellSize(1)
								, minP(2) + (dz + 0.5f) * cellSize(2);

					const auto cell = getCell(centerPoint, nullptr);
					for (size_t i = 1; i < 7; ++i)
						if ((cell.v[i]->distance >= 0 && cell.v[0]->distance < 0)
						    || (cell.v[i]->distance <= 0 && cell.v[0]->distance > 0)) {

							Vector3f gradient, centerPointExtendWithGradient;
							float centerPointDistance = distance(centerPoint, &gradient);
							gradient.normalize();
							centerPointExtendWithGradient = centerPoint + gradient;

							Info("Zero isocontour: p = %s, distance = %f, gradient = %s\n",
							     EigenUtility::toString(centerPoint).c_str(),
							     centerPointDistance,
							     EigenUtility::toString(gradient).c_str());

							contourFile << 'p'
							            << ' ' << centerPoint(0) << ' ' << centerPoint(1) << ' ' << centerPoint(2) << '\n';
							contourFile << 'l'
							            << ' ' << centerPoint(0) << ' ' << centerPoint(1) << ' ' << centerPoint(2)
							            << ' ' << centerPointExtendWithGradient(0) << ' ' << centerPointExtendWithGradient(1)
							            << ' ' << centerPointExtendWithGradient(2) << '\n';

							++successCounter;

							break;
					}
				}

		Info("Successful write %d points to the file \"%s\"\n", (int)successCounter, filePath.c_str());
		return true;
	}
}